package seq_diverse.main;
import org.jfree.data.xy.XYSeries;

public class DirectMap {
    int kmer = 0;
    long [][] counts;
    int perms = 0;
    int length;
    
    boolean done = false;
    XYSeries series0;
    XYSeries series1;
    XYSeries series2;
    XYSeries series3;

    public DirectMap(int len, int k){
	length = len;
	kmer = k;
	int width = length - kmer + 1;
	perms = pow(4, kmer);
	counts = new long[width][perms];
    }

    public void insert(String read){
	int key = 0;
	int mask = perms - 1;
	int ns = 0;

	for(int i = 0; i < Math.min(read.length(), length); i++){
	    int bits = getBits(read.charAt(i));
	    if(bits == 4){
		ns = (ns << 2 | 3) & mask;
	    }else{
		ns = (ns << 2) & mask;
	    }

	    key = (key << 2 | (bits & 3)) & mask;
	    
	    if(i >= kmer -1){
		if(ns == 0){
		    done = false;
		    counts[i-kmer+1][key]++;
		}
	    }
	    
	}
	
    }

    public XYSeries getMaximum(){
	if(!done){
	    calculateSeries();
	}
	return series0;
    }

    public XYSeries getDiversity(){
	if(!done){
	    calculateSeries();
	}
	return series1;
    }
    
    public XYSeries getRichness(){
	if(!done){
	    calculateSeries();
	}
	return series2;
    }
    
    public XYSeries getEvenness(){
	if(!done){
	    calculateSeries();
	}
	return series3;
    }

    private int pow(int a, int b){
	int ans = 1;
	for(int i = 0; i < b; i++){
	    ans = ans*a;
	}
	return ans;
    }

    private void calculateSeries(){
	series0 = new XYSeries("Maximum");
	series1 = new XYSeries("Diversity");
	series2 = new XYSeries("Richness");
	series3 = new XYSeries("Evenness");

	System.out.println("Pos.\tk-mer Richness\tk-mer Evenness\tk-mer Diversity");
		
	for(int i = 0; i < counts.length; i++){
	    double unique = 0;
	    double total = 0;
	    double xlogx = 0.0;

	    for(int j = 0; j < counts[i].length; j++){
		double x = counts[i][j];
		if(x > 0){
		    unique++;
		    total += x;
		    xlogx += x*Math.log(x);
		}
	    }
	    if(total == 0){
		total = 1;
		unique = 1;
	    }

	    double diversity = Math.log(total) - xlogx/total;
	    
	    double rich = 0.0;
	    double even = 0.0;
	    double dive = 0.0;

	    series0.add(i+1, 1.0);
	    
	    dive = Math.exp(diversity)/pow(4, kmer);
	    series1.add(i+1, dive);

	    if(unique < 1){
		rich = 0;
		even = 0;
	    }else{
		rich = unique/pow(4, kmer);
		even = Math.exp(diversity)/unique;
	    }

	    series2.add(i+1, rich);
	    series3.add(i+1, even);
	    
	    System.out.print((i+1) + "\t");
	    System.out.print(rich + "\t");
	    System.out.print(even + "\t");
	    System.out.println(dive);
	}

	done = true;
    }
    
    private int getBits(char c){
	if(c == 'a' || c == 'A') return 0;
	if(c == 'c' || c == 'C') return 1;
	if(c == 'g' || c == 'G') return 2;
	if(c == 't' || c == 'T') return 3;
	return 4;
    }
}